package com.project.demo.util;

import com.project.demo.attribute.Attribute;
import com.project.demo.engine.CpAneEngine;
import com.project.demo.entity.SysKey;
import com.project.demo.parameter.PublicKey;
import com.project.demo.parameter.SystemKey;
import com.project.demo.parameter.UserPrivateKey;
import com.project.demo.service.UserService;
import com.project.demo.structure.AccessTree;
import com.project.demo.structure.AccessTreeBuildModel;
import com.project.demo.text.CipherText;
import com.project.demo.text.PlainText;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;

public class ABEEncryption {


    public static AccessTree getAccessTree(PublicKey publicKey) {
        AccessTreeBuildModel[] accessTreeBuildModels = new AccessTreeBuildModel[7];
        //根节点ID必须为1
        accessTreeBuildModels[0] = AccessTreeBuildModel.innerAccessTreeBuildModel(1, 2, 1, -1);
        accessTreeBuildModels[1] = AccessTreeBuildModel.leafAccessTreeBuildModel(2, 1, "医生", 1);
        accessTreeBuildModels[2] = AccessTreeBuildModel.leafAccessTreeBuildModel(3, 2, "护士", 1);
        accessTreeBuildModels[3] = AccessTreeBuildModel.leafAccessTreeBuildModel(4, 3, "管理员", 1);
        accessTreeBuildModels[4] = AccessTreeBuildModel.innerAccessTreeBuildModel(5, 1, 4, 1);
        accessTreeBuildModels[5] = AccessTreeBuildModel.leafAccessTreeBuildModel(6, 1, "患者", 5);
        accessTreeBuildModels[6] = AccessTreeBuildModel.leafAccessTreeBuildModel(7, 2, "游客", 5);
        return AccessTree.build(publicKey, accessTreeBuildModels);
    }

    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/ECB/PKCS5Padding";

    public static byte[] encrypt(byte[] data, String key) throws Exception {
        SysKey sysKey = UserService.map.get("SystemKey");
        SystemKey systemKey = new SystemKey();
        String plainTextStr = new String(data);
        PlainText plainText = new PlainText(plainTextStr, systemKey.getPublicKey());
        //5.构建访问树
        systemKey =SystemKey.build();
        AccessTree accessTree = getAccessTree(systemKey.getPublicKey());
        //6.加密
        CpAneEngine cpAneEngine = new CpAneEngine();
        cpAneEngine.encrypt(systemKey.getPublicKey(), plainText, accessTree);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(), ALGORITHM);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
        return cipher.doFinal(data);
    }

    public static byte[] decrypt(byte[] encryptedData, String key) throws Exception {
        SysKey sysKey = UserService.map.get("SystemKey");
        SystemKey systemKey = new SystemKey();
        String plainTextStr = new String(encryptedData);
        CpAneEngine cpAneEngine = new CpAneEngine();
        //2.设置用户属性
        List<Attribute> attributes = Arrays.asList(
                new Attribute(plainTextStr, systemKey.getPublicKey())
        );
        systemKey = SystemKey.build();
        PlainText plainText = new PlainText(plainTextStr, systemKey.getPublicKey());
        UserPrivateKey userPrivateKey = cpAneEngine.keyGen(systemKey.getMasterPrivateKey(), attributes);
        CipherText cipherText  = cpAneEngine.encrypt(systemKey.getPublicKey(), plainText, getAccessTree(systemKey.getPublicKey()));
        cpAneEngine.decryptToStr(systemKey.getPublicKey(), userPrivateKey, cipherText);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(), ALGORITHM);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
        return cipher.doFinal(encryptedData);
    }

    public static void encryptFile(File inputFile, File outputFile, String key) throws Exception {
        try (FileInputStream fis = new FileInputStream(inputFile);
             FileOutputStream fos = new FileOutputStream(outputFile)) {
            byte[] buffer = new byte[(int) inputFile.length()];
            fis.read(buffer);
            byte[] encryptedBytes = encrypt(buffer, key);
            fos.write(encryptedBytes);
        }
    }

    public static void decryptFile(File inputFile, File outputFile, String key) throws Exception {
        try (FileInputStream fis = new FileInputStream(inputFile);
             FileOutputStream fos = new FileOutputStream(outputFile)) {
            byte[] buffer = new byte[(int) inputFile.length()];
            fis.read(buffer);
            byte[] decryptedBytes = decrypt(buffer, key);
            fos.write(decryptedBytes);
        }
    }
}
